import numpy as np
import os


class SortProcess(object):
    def __init__(self, src_path, trg_path):
        self.src_path = src_path
        self.trg_path = trg_path

        self.pad_flag = '<PAD>'
        self.unk_flag = '<UNK>'
        self.beg_flag = '<GO>'
        self.end_flag = '<EOS>'

        self.chars = [self.pad_flag, self.unk_flag, self.beg_flag, self.end_flag] +\
                     [chr(i) for i in range(97, 97+26)]
        self.c2i = {c: i for i, c in enumerate(self.chars)}
        self.i2c = {i: c for i, c in enumerate(self.chars)}

        self.pad_id = self.c2i.get(self.pad_flag, 0)
        self.unk_id = self.c2i.get(self.unk_flag, 0)
        self.beg_id = self.c2i.get(self.beg_flag, 0)
        self.end_id = self.c2i.get(self.end_flag, 0)

    def read_files(self, path):
        with open(path, 'r') as f:
            data_s = f.read().split('\n')
        data_s = [l.strip() for l in data_s if len(l.strip()) > 0]
        return data_s

    def to_index(self, chars, max_len=10):
        chars = chars[:max_len]
        ids = [self.c2i.get(c, self.unk_id) for c in chars] + [self.pad_id]*(max_len-len(chars))
        return ids

    def dataset(self, max_len=10, to_numpy=True):
        src_texts = self.read_files(self.src_path)
        trg_texts = self.read_files(self.trg_path)

        src_ids, trg_ids, label_ids = [], [], []
        for s, t in zip(src_texts, trg_texts):
            s = s[: max_len-2]
            t = t[: max_len-2]
            src_ids.append(self.to_index(chars=[self.beg_flag] + list(s) +\
                                               [self.end_flag], max_len=max_len))
            trg_ids.append(self.to_index(chars=[self.beg_flag] + list(t), max_len=max_len))
            label_ids.append(self.to_index(chars=list(t) + [self.end_flag], max_len=max_len))
        if to_numpy:
            src_ids, trg_ids, label_ids = np.array(src_ids), np.array(trg_ids), np.array(label_ids)
        return src_ids, trg_ids, label_ids










